"""
Fortran call graph analyzer for CodeMCP framework.

This module provides call graph analysis for Fortran programs,
tracking procedure calls, function references, and control flow.
"""

import logging
from typing import Dict, List, Any, Set, Optional, Tuple
from pathlib import Path
from collections import defaultdict

try:
    from ...core.error_handler import AnalysisError, log_info
except ImportError:
    import sys
    from pathlib import Path
    sys.path.insert(0, str(Path(__file__).parent.parent.parent))
    from core.error_handler import AnalysisError, log_info


class CallNode:
    """Represents a node in the call graph."""
    
    def __init__(self, name: str, node_type: str, file_path: str = "", line_number: int = 0):
        self.name = name
        self.type = node_type  # 'program', 'module', 'function', 'subroutine', 'interface'
        self.file_path = file_path
        self.line_number = line_number
        self.calls = set()  # Set of called procedures
        self.called_by = set()  # Set of procedures that call this one
        self.arguments = []
        self.return_type = None
        self.attributes = {}  # Additional attributes (pure, elemental, recursive, etc.)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert node to dictionary representation."""
        return {
            'name': self.name,
            'type': self.type,
            'file_path': self.file_path,
            'line_number': self.line_number,
            'calls': list(self.calls),
            'called_by': list(self.called_by),
            'arguments': self.arguments,
            'return_type': self.return_type,
            'attributes': self.attributes
        }


class CallGraphAnalyzer:
    """
    Fortran call graph analyzer.
    
    Analyzes procedure calls and builds call graphs for Fortran programs,
    including support for:
    - CALL statements and function references
    - Module procedures and USE statements
    - Interface blocks and generic procedures
    - Recursive procedures
    - Internal procedures and host association
    """
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        self.nodes = {}  # name -> CallNode
        self.call_graph = {}  # caller -> set of callees
        self.reverse_graph = {}  # callee -> set of callers
        self.module_procedures = {}  # module -> set of procedures
        self.procedure_modules = {}  # procedure -> module
        self.interfaces = {}  # interface name -> list of procedures
        self._reset()
    
    def _reset(self):
        """Reset analyzer state."""
        self.nodes.clear()
        self.call_graph.clear()
        self.reverse_graph.clear()
        self.module_procedures.clear()
        self.procedure_modules.clear()
        self.interfaces.clear()
        self.current_file = ""
        self.current_module = None
        self.current_procedure = None
    
    def analyze_file(self, parsed_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Analyze call graph for a single file.
        
        Args:
            parsed_data: Parsed Fortran file data
            
        Returns:
            Call graph analysis results
        """
        try:
            filename = parsed_data.get('filename', 'unknown')
            log_info(f"Analyzing call graph for {filename}")
            
            self.current_file = filename
            
            # Build nodes from parsed data
            self._build_nodes_from_parsed_data(parsed_data)
            
            # Analyze calls and build graph
            self._analyze_calls(parsed_data)
            
            # Resolve interface calls
            self._resolve_interface_calls(parsed_data)
            
            # Compute graph metrics
            metrics = self._compute_graph_metrics()
            
            result = {
                'filename': filename,
                'nodes': {name: node.to_dict() for name, node in self.nodes.items()},
                'call_graph': {caller: list(callees) for caller, callees in self.call_graph.items()},
                'reverse_graph': {callee: list(callers) for callee, callers in self.reverse_graph.items()},
                'metrics': metrics,
                'cycles': self._detect_cycles(),
                'entry_points': self._find_entry_points(),
                'leaf_procedures': self._find_leaf_procedures()
            }
            
            log_info(f"Call graph analysis complete for {filename}: "
                    f"{len(self.nodes)} nodes, {sum(len(calls) for calls in self.call_graph.values())} edges")
            
            return result
            
        except Exception as e:
            self.logger.error(f"Error analyzing call graph: {e}")
            raise AnalysisError(f"Call graph analysis failed: {e}")
    
    def _build_nodes_from_parsed_data(self, parsed_data: Dict[str, Any]):
        """Build call graph nodes from parsed data."""
        filename = parsed_data.get('filename', 'unknown')
        
        # Add program units
        for program in parsed_data.get('programs', []):
            node = CallNode(
                name=program['name'],
                node_type='program',
                file_path=filename,
                line_number=program.get('line_number', 0)
            )
            self.nodes[program['name']] = node
        
        # Add modules
        for module in parsed_data.get('modules', []):
            node = CallNode(
                name=module['name'],
                node_type='module',
                file_path=filename,
                line_number=module.get('line_number', 0)
            )
            self.nodes[module['name']] = node
            self.module_procedures[module['name']] = set(module.get('procedures', []))
            
            # Mark module procedures
            for proc_name in module.get('procedures', []):
                self.procedure_modules[proc_name] = module['name']
        
        # Add functions
        for function in parsed_data.get('functions', []):
            node = CallNode(
                name=function['name'],
                node_type='function',
                file_path=filename,
                line_number=function.get('line_number', 0)
            )
            node.return_type = function.get('return_type', 'unknown')
            node.arguments = function.get('arguments', [])
            node.attributes = {
                'is_recursive': function.get('is_recursive', False),
                'is_pure': function.get('is_pure', False),
                'is_elemental': function.get('is_elemental', False)
            }
            self.nodes[function['name']] = node
        
        # Add subroutines
        for subroutine in parsed_data.get('subroutines', []):
            node = CallNode(
                name=subroutine['name'],
                node_type='subroutine',
                file_path=filename,
                line_number=subroutine.get('line_number', 0)
            )
            node.arguments = subroutine.get('arguments', [])
            node.attributes = {
                'is_recursive': subroutine.get('is_recursive', False),
                'is_pure': subroutine.get('is_pure', False),
                'is_elemental': subroutine.get('is_elemental', False)
            }
            self.nodes[subroutine['name']] = node
        
        # Add interfaces
        for interface in parsed_data.get('interfaces', []):
            node = CallNode(
                name=interface['name'],
                node_type='interface',
                file_path=filename,
                line_number=interface.get('line_number', 0)
            )
            node.attributes = {
                'is_generic': interface.get('is_generic', True),
                'is_operator': interface.get('is_operator', False),
                'is_assignment': interface.get('is_assignment', False)
            }
            self.nodes[interface['name']] = node
            self.interfaces[interface['name']] = interface.get('procedures', [])
    
    def _analyze_calls(self, parsed_data: Dict[str, Any]):
        """Analyze procedure calls and build call graph edges."""
        
        # Process explicit calls from parsed data
        for call in parsed_data.get('calls', []):
            caller = call.get('in_procedure', 'unknown')
            callee = call.get('target', 'unknown')
            
            if caller != 'unknown' and callee != 'unknown':
                self._add_call_edge(caller, callee)
        
        # Process calls found in functions and subroutines
        for function in parsed_data.get('functions', []):
            caller = function['name']
            for callee in function.get('calls', []):
                self._add_call_edge(caller, callee)
        
        for subroutine in parsed_data.get('subroutines', []):
            caller = subroutine['name']
            for callee in subroutine.get('calls', []):
                self._add_call_edge(caller, callee)
    
    def _add_call_edge(self, caller: str, callee: str):
        """Add a call edge to the graph."""
        # Initialize graph structures
        if caller not in self.call_graph:
            self.call_graph[caller] = set()
        if callee not in self.reverse_graph:
            self.reverse_graph[callee] = set()
        
        # Add edge
        self.call_graph[caller].add(callee)
        self.reverse_graph[callee].add(caller)
        
        # Update node references
        if caller in self.nodes:
            self.nodes[caller].calls.add(callee)
        if callee in self.nodes:
            self.nodes[callee].called_by.add(caller)
    
    def _resolve_interface_calls(self, parsed_data: Dict[str, Any]):
        """Resolve calls through interface blocks."""
        for interface_name, procedures in self.interfaces.items():
            # Find calls to this interface
            for caller, callees in self.call_graph.items():
                if interface_name in callees:
                    # Replace interface call with calls to specific procedures
                    callees.remove(interface_name)
                    for proc in procedures:
                        callees.add(proc)
                        self._add_call_edge(caller, proc)
    
    def _compute_graph_metrics(self) -> Dict[str, Any]:
        """Compute call graph metrics."""
        total_nodes = len(self.nodes)
        total_edges = sum(len(callees) for callees in self.call_graph.values())
        
        # Compute in-degree and out-degree statistics
        in_degrees = [len(callers) for callers in self.reverse_graph.values()]
        out_degrees = [len(callees) for callees in self.call_graph.values()]
        
        # Find strongly connected components (simplified)
        cycles = self._detect_cycles()
        
        return {
            'total_nodes': total_nodes,
            'total_edges': total_edges,
            'average_in_degree': sum(in_degrees) / len(in_degrees) if in_degrees else 0,
            'average_out_degree': sum(out_degrees) / len(out_degrees) if out_degrees else 0,
            'max_in_degree': max(in_degrees) if in_degrees else 0,
            'max_out_degree': max(out_degrees) if out_degrees else 0,
            'cycles_detected': len(cycles),
            'is_dag': len(cycles) == 0
        }
    
    def _detect_cycles(self) -> List[List[str]]:
        """Detect cycles in the call graph using DFS."""
        cycles = []
        visited = set()
        rec_stack = set()
        path = []
        
        def dfs(node: str) -> bool:
            if node in rec_stack:
                # Found a cycle
                cycle_start = path.index(node)
                cycle = path[cycle_start:] + [node]
                cycles.append(cycle)
                return True
            
            if node in visited:
                return False
            
            visited.add(node)
            rec_stack.add(node)
            path.append(node)
            
            for neighbor in self.call_graph.get(node, []):
                if dfs(neighbor):
                    return True
            
            rec_stack.remove(node)
            path.pop()
            return False
        
        for node in self.nodes:
            if node not in visited:
                dfs(node)
        
        return cycles
    
    def _find_entry_points(self) -> List[str]:
        """Find entry points (nodes with no incoming calls)."""
        entry_points = []
        for node_name in self.nodes:
            if node_name not in self.reverse_graph or not self.reverse_graph[node_name]:
                # Check if it's a main program or public module procedure
                node = self.nodes[node_name]
                if (node.type == 'program' or 
                    (node.type in ['function', 'subroutine'] and 
                     node_name not in self.procedure_modules)):
                    entry_points.append(node_name)
        return entry_points
    
    def _find_leaf_procedures(self) -> List[str]:
        """Find leaf procedures (nodes with no outgoing calls)."""
        leaf_procedures = []
        for node_name in self.nodes:
            if node_name not in self.call_graph or not self.call_graph[node_name]:
                leaf_procedures.append(node_name)
        return leaf_procedures
    
    def build_call_tree(self, entry_point: str, max_depth: int = None) -> Dict[str, Any]:
        """
        Build a call tree starting from an entry point.
        
        Args:
            entry_point: Starting procedure name
            max_depth: Maximum depth to traverse (None for unlimited)
            
        Returns:
            Call tree representation
        """
        if entry_point not in self.nodes:
            raise AnalysisError(f"Entry point '{entry_point}' not found")
        
        visited = set()
        
        def build_tree(node_name: str, depth: int = 0) -> Dict[str, Any]:
            if max_depth is not None and depth > max_depth:
                return None
            
            node = self.nodes[node_name]
            tree_node = {
                'name': node_name,
                'type': node.type,
                'file_path': node.file_path,
                'line_number': node.line_number,
                'depth': depth,
                'children': []
            }
            
            # Avoid infinite recursion
            if node_name in visited:
                tree_node['is_recursive_call'] = True
                return tree_node
            
            visited.add(node_name)
            
            # Add children
            for callee in sorted(self.call_graph.get(node_name, [])):
                if callee in self.nodes:
                    child = build_tree(callee, depth + 1)
                    if child:
                        tree_node['children'].append(child)
            
            visited.remove(node_name)
            return tree_node
        
        return build_tree(entry_point)
    
    def get_call_path(self, from_proc: str, to_proc: str) -> List[List[str]]:
        """
        Find all call paths from one procedure to another.
        
        Args:
            from_proc: Source procedure
            to_proc: Target procedure
            
        Returns:
            List of call paths (each path is a list of procedure names)
        """
        if from_proc not in self.nodes or to_proc not in self.nodes:
            return []
        
        paths = []
        visited = set()
        
        def find_paths(current: str, target: str, path: List[str]):
            if current == target:
                paths.append(path[:])
                return
            
            if current in visited:
                return
            
            visited.add(current)
            
            for neighbor in self.call_graph.get(current, []):
                path.append(neighbor)
                find_paths(neighbor, target, path)
                path.pop()
            
            visited.remove(current)
        
        find_paths(from_proc, to_proc, [from_proc])
        return paths
    
    def get_procedure_dependencies(self, procedure: str) -> Dict[str, Any]:
        """
        Get all dependencies for a procedure.
        
        Args:
            procedure: Procedure name
            
        Returns:
            Dictionary with direct and transitive dependencies
        """
        if procedure not in self.nodes:
            return {}
        
        direct_deps = set(self.call_graph.get(procedure, []))
        
        # Find transitive dependencies
        transitive_deps = set()
        visited = set()
        
        def collect_deps(proc: str):
            if proc in visited:
                return
            visited.add(proc)
            
            for dep in self.call_graph.get(proc, []):
                transitive_deps.add(dep)
                collect_deps(dep)
        
        for dep in direct_deps:
            collect_deps(dep)
        
        return {
            'direct_dependencies': list(direct_deps),
            'transitive_dependencies': list(transitive_deps),
            'total_dependencies': len(transitive_deps)
        }